package com.demo.easyexcel.controller;

import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.write.metadata.WriteSheet;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.demo.easyexcel.entity.User;
import com.demo.easyexcel.mapper.UserMapper;
import org.springframework.util.CollectionUtils;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.OutputStream;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * 用户表(User)表控制层
 */
@RestController
@RequestMapping("user")
public class UserController {

    ThreadPoolExecutor executor = new ThreadPoolExecutor(20,
            40,
            50,
            TimeUnit.SECONDS,
            new LinkedBlockingQueue<>(20),
            new ThreadPoolExecutor.CallerRunsPolicy());

    @Resource
    private UserMapper userMapper;


    /**
     * 多线程导出Excel
     */
    @RequestMapping("/download")
    public void download(HttpServletResponse response) throws IOException {
        System.out.println("---------多线程查询导出开始-----------");
        long start = System.currentTimeMillis();
        OutputStream outputStream = response.getOutputStream();
        // 设置响应结果
        setExcelResponse(response);
        try {
            // 获取总数据量
            Long count = userMapper.selectCount(null);
            // 每页多少条数据
            int pageSize = 100000;
            ExcelWriter excelWriter = EasyExcel.write(outputStream).build();
            // 创建CompletableFuture集合
            List<CompletableFuture<Void>> completableFutures = new ArrayList<>();
            for (int i = 0; i < (count / pageSize) + 1; i++) {
                int index = i;
                // 异步执行查询数据方法
                CompletableFuture<Void> completableFuture = CompletableFuture.runAsync(() -> {
                    // 分页查询数据
                    Map<String, Object> map = new HashMap();
                    map.put("current", index * pageSize);
                    map.put("size", pageSize);
                    List<User> list = userMapper.selectUser(map);
                    if (!CollectionUtils.isEmpty(list)) {
                        WriteSheet writeSheet = EasyExcel.writerSheet(index, "user" + (index + 1)).head(User.class).build();
                        // 并发写入sheet会抛出异常
                        synchronized (excelWriter) {
                            excelWriter.write(list, writeSheet);
                        }
                    }
                }, executor);
                completableFutures.add(completableFuture);
            }
            // 等待线程执行完毕
            for (CompletableFuture<Void> completableFuture : completableFutures) {
                completableFuture.join();
            }
            // 整个excel写完关闭流
            excelWriter.finish();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        outputStream.flush();
        response.getOutputStream().close();
        System.out.println("共消耗：  " + (System.currentTimeMillis() - start) + "ms");
    }

    /**
     * 单线程导出Excel
     */
    @RequestMapping("/download1")
    public void download1(HttpServletResponse response) throws IOException {
        System.out.println("-----------单线程导出开始-----------");
        // 总共耗时
        long start = System.currentTimeMillis();
        // 设置响应结果
        setExcelResponse(response);
        // 查询导出数据总条数
        Long count = userMapper.selectCount(null);
        // 每个sheet页10W条数据
        int pageSize = 100000;
        OutputStream outputStream = response.getOutputStream();
        ExcelWriter excelWriter = EasyExcel.write(outputStream).build();
        for (int i = 0; i < (count / pageSize) + 1; i++) {
            // MybatisPlus分页查询数据
            Page<User> userIPage = userMapper.selectPage(new Page<>(i + 1, pageSize), null);
            List<User> records = userIPage.getRecords();
            WriteSheet writeSheet = EasyExcel.writerSheet(i, "user" + (i + 1)).head(User.class).build();
            excelWriter.write(records, writeSheet);

        }
        excelWriter.finish();
        outputStream.flush();
        outputStream.close();
        System.out.println("共消耗：  " + (System.currentTimeMillis() - start) + "ms");
    }

    private void setExcelResponse(HttpServletResponse response) throws UnsupportedEncodingException {
        // 设置响应内容
        response.setContentType("application/vnd.ms-excel");
        // 设置文件编码
        response.setCharacterEncoding("UTF-8");
        // 文件以附件形式下载
        response.setHeader("Content-disposition", "attachment;filename=demo_" + URLEncoder.encode(System.currentTimeMillis() + ".xlsx", "utf-8"));
    }

}

